





<!DOCTYPE html>
<html class="writer-html5" lang="zh-CN" >
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>向 Realy 中添加一个算子 &mdash; tvm 0.8.dev1982 文档</title>
  

  
  <link rel="stylesheet" href="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous">
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/gallery.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/tlcpack_theme.css" type="text/css" />

  
  
    <link rel="shortcut icon" href="../../_static/tvm-logo-square.png"/>
  

  
  
  
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
        <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script>
        <script src="../../_static/jquery.js"></script>
        <script src="../../_static/underscore.js"></script>
        <script src="../../_static/doctools.js"></script>
        <script src="../../_static/translations.js"></script>
        <script async="async" src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
    
    <script type="text/javascript" src="../../_static/js/theme.js"></script>

    
    <script type="text/javascript" src="../../_static/js/tlcpack_theme.js"></script>
    <link rel="index" title="索引" href="../../genindex.html" />
    <link rel="search" title="搜索" href="../../search.html" />
    <link rel="next" title="Adding a Compiler Pass to Relay" href="relay_add_pass.html" />
    <link rel="prev" title="开发者指南" href="how_to.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    
<header class="header">
    <div class="innercontainer">
      <div class="headerInner d-flex justify-content-between align-items-center">
          <div class="headerLogo">
               <a href="https://tvm.apache.org/"><img src=https://tvm.apache.org/assets/images/logo.svg alt="logo"></a>
          </div>

          <div id="headMenu" class="headerNav">
            <button type="button" id="closeHeadMenu" class="navCloseBtn"><img src="../../_static/img/close-icon.svg" alt="Close"></button>
             <ul class="nav">
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/community>Community</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/download>Download</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/vta>VTA</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/blog>Blog</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/docs>Docs</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmconf.org>Conference</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://github.com/apache/tvm/>Github</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmchinese.github.io/declaration_zh_CN.html>About-Translators</a>
                </li>
             </ul>
               <div class="responsivetlcdropdown">
                 <button type="button" class="btn-link">
                   ASF
                 </button>
                 <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                 </ul>
               </div>
          </div>
            <div class="responsiveMenuIcon">
              <button type="button" id="menuBtn" class="btn-menu"><img src="../../_static/img/menu-icon.svg" alt="Menu Icon"></button>
            </div>

            <div class="tlcDropdown">
              <div class="dropdown">
                <button type="button" class="btn-link dropdown-toggle" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">
                  ASF
                </button>
                <div class="dropdown-menu dropdown-menu-right">
                  <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                  </ul>
                </div>
              </div>
          </div>
       </div>
    </div>
 </header>
 
    <nav data-toggle="wy-nav-shift" class="wy-nav-side fixed">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../index.html">
          

          
            
            <img src="../../_static/tvm-logo-small.png" class="logo" alt="Logo"/>
          
          </a>

          
            
            
                <div class="version">
                  0.8.dev1982
                </div>
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption" role="heading"><span class="caption-text">如何开始</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../install/index.html">安装 TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../contribute/index.html">贡献者指南</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">用户引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../tutorial/index.html">User Tutorial</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../how_to/index.html">How To Guides</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">开发者引导</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="../tutorial/index.html">Developer Tutorial</a></li>
<li class="toctree-l1 current"><a class="reference internal" href="how_to.html">开发者指南</a><ul class="current">
<li class="toctree-l2 current"><a class="current reference internal" href="#">向 Realy 中添加一个算子</a><ul>
<li class="toctree-l3"><a class="reference internal" href="#defining-an-attribute-node">1. Defining an Attribute Node</a></li>
<li class="toctree-l3"><a class="reference internal" href="#writing-a-type-relation">2. Writing a Type Relation</a></li>
<li class="toctree-l3"><a class="reference internal" href="#relating-the-arity-and-attributes-to-an-operation">3. Relating the Arity and Attributes to an Operation</a></li>
<li class="toctree-l3"><a class="reference internal" href="#defining-the-compute-of-the-operation">4. Defining the Compute of the Operation</a></li>
<li class="toctree-l3"><a class="reference internal" href="#hooking-up-compute-and-strategy-with-relay">5. Hooking up Compute and Strategy with Relay</a></li>
<li class="toctree-l3"><a class="reference internal" href="#creating-a-relay-call-node-and-exposing-a-python-hook">6. Creating a Relay Call Node and Exposing a Python Hook</a></li>
<li class="toctree-l3"><a class="reference internal" href="#including-a-cleaner-python-api-hook">7. Including a Cleaner Python API Hook</a></li>
<li class="toctree-l3"><a class="reference internal" href="#writing-unit-tests">8. Writing Unit Tests!</a></li>
<li class="toctree-l3"><a class="reference internal" href="#other-topics">Other Topics</a><ul>
<li class="toctree-l4"><a class="reference internal" href="#gradient-operators">Gradient Operators</a></li>
<li class="toctree-l4"><a class="reference internal" href="#adding-a-gradient-in-python">Adding a Gradient in Python</a></li>
<li class="toctree-l4"><a class="reference internal" href="#adding-a-gradient-in-c">Adding a Gradient in C++</a></li>
</ul>
</li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="relay_add_pass.html">Adding a Compiler Pass to Relay</a></li>
<li class="toctree-l2"><a class="reference internal" href="relay_bring_your_own_codegen.html">Bring Your Own Codegen To TVM</a></li>
<li class="toctree-l2"><a class="reference internal" href="pytest_target_parametrization.html">Python Target Parametrization</a></li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">架构指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../arch/index.html">Design and Architecture</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">主题引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../topic/microtvm/index.html">microTVM：裸机使用TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../topic/vta/index.html">VTA: Versatile Tensor Accelerator</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">参考指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../reference/langref/index.html">语言参考</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/python/index.html">Python API</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/links.html">Other APIs</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/publications.html">Publications</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../genindex.html">索引</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
      
      <nav class="wy-nav-top" aria-label="top navigation" data-toggle="wy-nav-top">
        
            <div class="togglemenu">

            </div>
            <div class="nav-content">
              <!-- tvm -->
              Table of content
            </div>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        

          




















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../index.html">Docs</a> <span class="br-arrow">></span></li>
        
          <li><a href="how_to.html">开发者指南</a> <span class="br-arrow">></span></li>
        
      <li>向 Realy 中添加一个算子</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
            
            <a href="../../_sources/dev/how_to/relay_add_op.rst.txt" rel="nofollow"> <img src="../../_static//img/source.svg" alt="viewsource"/></a>
          
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <div class="section" id="adding-an-operator-to-relay">
<span id="relay-add-op"></span><h1>向 Realy 中添加一个算子<a class="headerlink" href="#adding-an-operator-to-relay" title="永久链接至标题">¶</a></h1>
<p>In this document we will go over the steps needed to register a new TVM operator
in Relay. We will be following this PR which adds a <a class="reference external" href="https://github.com/apache/tvm/pull/7722">cumulative product</a> operation as an example.
The PR itself builds upon another PR which adds a <a class="reference external" href="https://github.com/apache/tvm/pull/7334">cumulative sum</a> operation.</p>
<p>Registering a new operator requires a few steps:</p>
<ol class="arabic simple">
<li><p>Add an attribute node declaring fixed arguments which are known at compile time</p></li>
<li><p>Write a type relation for your operation to integrate into Relay’s type system.</p></li>
<li><p>Use the <code class="docutils literal notranslate"><span class="pre">RELAY_REGISTER_OP</span></code> macro in C++ to register the operator’s arity, type, and other hints for the compiler</p></li>
<li><p>Write how the operator is computed</p></li>
<li><p>Register the compute, schedule with the relay operator</p></li>
<li><p>Define a C++ function to produce a call node for the operator and registering a Python API hook for the function</p></li>
<li><p>Wrapping the above Python API hook in a neater interface</p></li>
<li><p>Writing tests for the new relay operator</p></li>
</ol>
<div class="section" id="defining-an-attribute-node">
<h2>1. Defining an Attribute Node<a class="headerlink" href="#defining-an-attribute-node" title="永久链接至标题">¶</a></h2>
<p>Attributes are fixed arguments which are supposed to be known at compile time. The stride and dilation of a convolution
operator would be an appropriate example of fields which might belong in an attribute node for a convolution operator.</p>
<p>Attributes should be defined in a file within the folder <a class="reference external" href="https://github.com/apache/tvm/tree/main/include/tvm/relay/attrs">include/tvm/relay/attrs/</a>.</p>
<p>Ultimately we want to create an operator whose interface can be seen clearly in the final python interface:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">cumprod</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="bp">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">None</span><span class="p">,</span> <span class="n">exclusive</span><span class="o">=</span><span class="bp">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Numpy style cumprod op. Return the cumulative inclusive product of the elements along</span>
<span class="sd">    a given axis.</span>
<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    data : relay.Expr</span>
<span class="sd">        The input data to the operator.</span>
<span class="sd">    axis : int, optional</span>
<span class="sd">        Axis along which the cumulative product is computed. The default (None) is to compute</span>
<span class="sd">        the cumprod over the flattened array.</span>
<span class="sd">    dtype : string, optional</span>
<span class="sd">        Type of the returned array and of the accumulator in which the elements are multiplied.</span>
<span class="sd">        If dtype is not specified, it defaults to the dtype of data.</span>
<span class="sd">    exclusive : bool, optional</span>
<span class="sd">        If true will return exclusive product in which the first element is not</span>
<span class="sd">        included. In other terms, if true, the j-th output element would be</span>
<span class="sd">        the product of the first (j-1) elements. Otherwise, it would be the product of</span>
<span class="sd">        the first j elements. The product of zero elements will be 1.</span>
<span class="sd">    Returns</span>
<span class="sd">    -------</span>
<span class="sd">    result : relay.Expr</span>
<span class="sd">        The result has the same size as data, and the same shape as data if axis is not None.</span>
<span class="sd">        If axis is None, the result is a 1-d array.</span>
<span class="sd">    &quot;&quot;&quot;</span>
</pre></div>
</div>
<p>A similiar interface exists for <code class="docutils literal notranslate"><span class="pre">cumsum()</span></code>.</p>
<p>Therefore, when defining our attributes in <code class="docutils literal notranslate"><span class="pre">include/tvm/relay/attrs/transform.h</span></code> we choose the axis,
accumulation dtype, and exclusivity of the operation as appropriate fields for the struct.</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="cm">/*! \brief Attributes used in cumsum and cumprod operator */</span>
<span class="k">struct</span> <span class="nl">ScanopAttrs</span> <span class="p">:</span> <span class="k">public</span> <span class="n">tvm</span><span class="o">::</span><span class="n">AttrsNode</span><span class="o">&lt;</span><span class="n">ScanopAttrs</span><span class="o">&gt;</span> <span class="p">{</span>
  <span class="n">Integer</span> <span class="n">axis</span><span class="p">;</span>
  <span class="n">DataType</span> <span class="n">dtype</span><span class="p">;</span>
  <span class="n">Bool</span> <span class="n">exclusive</span> <span class="o">=</span> <span class="n">Bool</span><span class="p">(</span><span class="nb">false</span><span class="p">);</span>
  <span class="n">TVM_DECLARE_ATTRS</span><span class="p">(</span><span class="n">ScanopAttrs</span><span class="p">,</span> <span class="s">&quot;relay.attrs.ScanopAttrs&quot;</span><span class="p">)</span> <span class="p">{</span>
    <span class="n">TVM_ATTR_FIELD</span><span class="p">(</span><span class="n">axis</span><span class="p">).</span><span class="n">describe</span><span class="p">(</span><span class="s">&quot;The axis to operate over&quot;</span><span class="p">).</span><span class="n">set_default</span><span class="p">(</span><span class="n">NullValue</span><span class="o">&lt;</span><span class="n">Integer</span><span class="o">&gt;</span><span class="p">());</span>
    <span class="n">TVM_ATTR_FIELD</span><span class="p">(</span><span class="n">dtype</span><span class="p">).</span><span class="n">describe</span><span class="p">(</span><span class="s">&quot;Output data type&quot;</span><span class="p">).</span><span class="n">set_default</span><span class="p">(</span><span class="n">NullValue</span><span class="o">&lt;</span><span class="n">DataType</span><span class="o">&gt;</span><span class="p">());</span>
    <span class="n">TVM_ATTR_FIELD</span><span class="p">(</span><span class="n">exclusive</span><span class="p">)</span>
        <span class="p">.</span><span class="n">describe</span><span class="p">(</span><span class="s">&quot;The first element is not included&quot;</span><span class="p">)</span>
        <span class="p">.</span><span class="n">set_default</span><span class="p">(</span><span class="n">Bool</span><span class="p">(</span><span class="nb">false</span><span class="p">));</span>
  <span class="p">}</span>
<span class="p">};</span>
</pre></div>
</div>
</div>
<div class="section" id="writing-a-type-relation">
<h2>2. Writing a Type Relation<a class="headerlink" href="#writing-a-type-relation" title="永久链接至标题">¶</a></h2>
<p>To allow for flexibility in registering operators and greater
expressivity and granularity in expressing types in Relay, operators
are typed using relations between input and output types. These relations
are represented as functions that take in a list of input types and
output types (any of these types may be incomplete) and return a list
of input and output types that satisfies the relation. This includes shape
information which can be determined statically at compile time. Essentially, a
relation for an operator can enforce all the necessary typing rules
(namely by inspecting the input types) in addition to computing the
output type.</p>
<p>Type relation for the cumulative product and sum operators can be found in
<code class="docutils literal notranslate"><span class="pre">src/relay/op/tensor/transform.cc</span></code>:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">TVM_REGISTER_NODE_TYPE</span><span class="p">(</span><span class="n">ScanopAttrs</span><span class="p">);</span>
<span class="kt">bool</span> <span class="nf">ScanopRel</span><span class="p">(</span><span class="k">const</span> <span class="n">Array</span><span class="o">&lt;</span><span class="n">Type</span><span class="o">&gt;&amp;</span> <span class="n">types</span><span class="p">,</span> <span class="kt">int</span> <span class="n">num_inputs</span><span class="p">,</span> <span class="k">const</span> <span class="n">Attrs</span><span class="o">&amp;</span> <span class="n">attrs</span><span class="p">,</span> <span class="k">const</span> <span class="n">TypeReporter</span><span class="o">&amp;</span> <span class="n">reporter</span><span class="p">)</span> <span class="p">{</span>
    <span class="c1">// types: [data, output]</span>
    <span class="n">ICHECK_EQ</span><span class="p">(</span><span class="n">types</span><span class="p">.</span><span class="n">size</span><span class="p">(),</span> <span class="mi">2</span><span class="p">)</span> <span class="o">&lt;&lt;</span> <span class="s">&quot;Expects two types, one for the input and another for the output&quot;</span><span class="p">;</span>
    <span class="k">const</span> <span class="k">auto</span><span class="o">*</span> <span class="n">data</span> <span class="o">=</span> <span class="n">types</span><span class="p">[</span><span class="mi">0</span><span class="p">].</span><span class="n">as</span><span class="o">&lt;</span><span class="n">TensorTypeNode</span><span class="o">&gt;</span><span class="p">();</span>
    <span class="k">if</span> <span class="p">(</span><span class="n">data</span> <span class="o">==</span> <span class="k">nullptr</span><span class="p">)</span> <span class="p">{</span>
        <span class="n">ICHECK</span><span class="p">(</span><span class="n">types</span><span class="p">[</span><span class="mi">0</span><span class="p">].</span><span class="n">as</span><span class="o">&lt;</span><span class="n">IncompleteTypeNode</span><span class="o">&gt;</span><span class="p">())</span>
        <span class="o">&lt;&lt;</span> <span class="s">&quot;Scanop: expect input type to be TensorType but get &quot;</span> <span class="o">&lt;&lt;</span> <span class="n">types</span><span class="p">[</span><span class="mi">0</span><span class="p">];</span>
        <span class="k">return</span> <span class="nb">false</span><span class="p">;</span>
    <span class="p">}</span>

    <span class="k">const</span> <span class="k">auto</span><span class="o">*</span> <span class="n">param</span> <span class="o">=</span> <span class="n">attrs</span><span class="p">.</span><span class="n">as</span><span class="o">&lt;</span><span class="n">ScanopAttrs</span><span class="o">&gt;</span><span class="p">();</span>

    <span class="k">auto</span> <span class="n">dtype</span> <span class="o">=</span> <span class="n">param</span><span class="o">-&gt;</span><span class="n">dtype</span><span class="p">;</span>
    <span class="k">if</span> <span class="p">(</span><span class="n">dtype</span><span class="p">.</span><span class="n">is_void</span><span class="p">())</span> <span class="p">{</span>
        <span class="n">dtype</span> <span class="o">=</span> <span class="n">data</span><span class="o">-&gt;</span><span class="n">dtype</span><span class="p">;</span>
    <span class="p">}</span>

    <span class="k">if</span> <span class="p">(</span><span class="n">param</span><span class="o">-&gt;</span><span class="n">axis</span><span class="p">.</span><span class="n">defined</span><span class="p">())</span> <span class="p">{</span>
        <span class="n">reporter</span><span class="o">-&gt;</span><span class="n">Assign</span><span class="p">(</span><span class="n">types</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">TensorType</span><span class="p">(</span><span class="n">data</span><span class="o">-&gt;</span><span class="n">shape</span><span class="p">,</span> <span class="n">dtype</span><span class="p">));</span>
    <span class="p">}</span> <span class="k">else</span> <span class="p">{</span>
        <span class="k">auto</span> <span class="n">prod</span> <span class="o">=</span> <span class="n">data</span><span class="o">-&gt;</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">];</span>
        <span class="k">for</span> <span class="p">(</span><span class="kt">size_t</span> <span class="n">i</span> <span class="o">=</span> <span class="mi">1</span><span class="p">;</span> <span class="n">i</span> <span class="o">&lt;</span> <span class="n">data</span><span class="o">-&gt;</span><span class="n">shape</span><span class="p">.</span><span class="n">size</span><span class="p">();</span> <span class="o">++</span><span class="n">i</span><span class="p">)</span> <span class="p">{</span>
            <span class="n">prod</span> <span class="o">=</span> <span class="n">prod</span> <span class="o">*</span> <span class="n">data</span><span class="o">-&gt;</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">];</span>
        <span class="p">}</span>
        <span class="n">reporter</span><span class="o">-&gt;</span><span class="n">Assign</span><span class="p">(</span><span class="n">types</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">TensorType</span><span class="p">({</span><span class="n">prod</span><span class="p">},</span> <span class="n">dtype</span><span class="p">));</span>
    <span class="p">}</span>

    <span class="k">return</span> <span class="nb">true</span><span class="p">;</span>
<span class="p">}</span>
</pre></div>
</div>
</div>
<div class="section" id="relating-the-arity-and-attributes-to-an-operation">
<h2>3. Relating the Arity and Attributes to an Operation<a class="headerlink" href="#relating-the-arity-and-attributes-to-an-operation" title="永久链接至标题">¶</a></h2>
<p>We then register the name of our new ops and annotate them with the calling interface.
The <code class="docutils literal notranslate"><span class="pre">RELAY_REGISTER_OP</span></code> macro in C++ allows a developer
to specify the following information about an operator in Relay:</p>
<ul class="simple">
<li><p>Arity (number of arguments)</p></li>
<li><p>Names and descriptions for positional arguments</p></li>
<li><p>Support level (1 indicates an internal intrinsic; higher numbers indicate less integral or externally supported operators)</p></li>
<li><p>A type relation for the operator</p></li>
<li><p>Other annotations useful when optimizing the operation.</p></li>
</ul>
<p>Once again we add this to <code class="docutils literal notranslate"><span class="pre">src/relay/op/tensor/transform.cc</span></code>:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">RELAY_REGISTER_OP</span><span class="p">(</span><span class="s">&quot;cumsum&quot;</span><span class="p">)</span>
    <span class="p">.</span><span class="n">describe</span><span class="p">(</span>
        <span class="sa">R</span><span class="s">&quot;</span><span class="dl">doc(</span><span class="s">Return the cumulative sum of the elements along a given axis.</span><span class="dl">)doc</span><span class="s">&quot;</span> <span class="n">TVM_ADD_FILELINE</span><span class="p">)</span>
    <span class="p">.</span><span class="n">set_num_inputs</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
    <span class="p">.</span><span class="n">add_argument</span><span class="p">(</span><span class="s">&quot;data&quot;</span><span class="p">,</span> <span class="s">&quot;Tensor&quot;</span><span class="p">,</span> <span class="s">&quot;The input tensor.&quot;</span><span class="p">)</span>
    <span class="p">.</span><span class="n">set_support_level</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
    <span class="p">.</span><span class="n">add_type_rel</span><span class="p">(</span><span class="s">&quot;Cumsum&quot;</span><span class="p">,</span> <span class="n">ScanopRel</span><span class="p">)</span>
    <span class="p">.</span><span class="n">set_attr</span><span class="o">&lt;</span><span class="n">TOpPattern</span><span class="o">&gt;</span><span class="p">(</span><span class="s">&quot;TOpPattern&quot;</span><span class="p">,</span> <span class="n">kOpaque</span><span class="p">);</span>

<span class="n">RELAY_REGISTER_OP</span><span class="p">(</span><span class="s">&quot;cumprod&quot;</span><span class="p">)</span>
    <span class="p">.</span><span class="n">describe</span><span class="p">(</span>
        <span class="sa">R</span><span class="s">&quot;</span><span class="dl">doc(</span><span class="s">Return the cumulative product of the elements along a given axis.</span><span class="dl">)doc</span><span class="s">&quot;</span> <span class="n">TVM_ADD_FILELINE</span><span class="p">)</span>
    <span class="p">.</span><span class="n">set_num_inputs</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
    <span class="p">.</span><span class="n">add_argument</span><span class="p">(</span><span class="s">&quot;data&quot;</span><span class="p">,</span> <span class="s">&quot;Tensor&quot;</span><span class="p">,</span> <span class="s">&quot;The input tensor.&quot;</span><span class="p">)</span>
    <span class="p">.</span><span class="n">set_support_level</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
    <span class="p">.</span><span class="n">add_type_rel</span><span class="p">(</span><span class="s">&quot;Cumprod&quot;</span><span class="p">,</span> <span class="n">ScanopRel</span><span class="p">)</span>
    <span class="p">.</span><span class="n">set_attr</span><span class="o">&lt;</span><span class="n">TOpPattern</span><span class="o">&gt;</span><span class="p">(</span><span class="s">&quot;TOpPattern&quot;</span><span class="p">,</span> <span class="n">kOpaque</span><span class="p">);</span>
</pre></div>
</div>
<p>In this case the <code class="docutils literal notranslate"><span class="pre">TOpPattern</span></code> is a hint to the compiler on the pattern of computation the operator does, which might be
useful for fusing operators. <code class="docutils literal notranslate"><span class="pre">kOpaque</span></code> tells TVM to not bother trying to fuse this operator.</p>
</div>
<div class="section" id="defining-the-compute-of-the-operation">
<h2>4. Defining the Compute of the Operation<a class="headerlink" href="#defining-the-compute-of-the-operation" title="永久链接至标题">¶</a></h2>
<p>While we’ve now defined the interface for our operations we still need to define
how to perform the actual calculations for cumulative sum and product.</p>
<p>Writing this code is outside the scope of the tutorial. For now, we assume
we have a well tested implementation for the operation’s compute. For
more details on how to do this, we recommend looking up the tutorials
on <a class="reference external" href="https://tvm.apache.org/docs/tutorials/get_started/tensor_expr_get_started.html">tensor expressions</a>, <a class="reference external" href="https://tvm.apache.org/docs/tutorials/topi/intro_topi.html">TVM’s operator inventory (topi)</a> and looking at the
example cumulative sum and product implementations found in <a class="reference external" href="https://github.com/apache/tvm/blob/main/python/tvm/topi/scan.py">python/tvm/topi/scan.py</a>
and the gpu versions in <a class="reference external" href="https://github.com/apache/tvm/blob/main/python/tvm/topi/cuda/scan.py">python/tvm/topi/cuda/scan.py</a>. In the case of our cumulative
sum and product operations we write things directly in <a class="reference external" href="https://tvm.apache.org/docs/dev/index.html?highlight=tir#tvm-tir">TIR</a> which is the
representation where tensor expressions and topi will lower into.</p>
</div>
<div class="section" id="hooking-up-compute-and-strategy-with-relay">
<h2>5. Hooking up Compute and Strategy with Relay<a class="headerlink" href="#hooking-up-compute-and-strategy-with-relay" title="永久链接至标题">¶</a></h2>
<p>After you have implemented your compute function we now need to glue it to our
relay operation. Within TVM this means not only defining the computation, but also the schedule
for an operation. A strategy is a method which picks which computation and which schedule
to use. For example, for 2D convolutions we might recognize we are doing a depthwise convolution
and dispatch to a more efficient computation and schedule as a result. In our case however we have
no such need except for dispatching between our CPU and GPU implementations. In
<code class="docutils literal notranslate"><span class="pre">python/tvm/relay/op/strategy/generic.py</span></code> and <code class="docutils literal notranslate"><span class="pre">python/tvm/relay/op/strategy/cuda.py</span></code> we
add the following strategies:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">wrap_compute_scanop</span><span class="p">(</span><span class="n">topi_compute</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Wrap scanop style topi compute&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">_compute_scanop</span><span class="p">(</span><span class="n">attrs</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">_</span><span class="p">):</span>
        <span class="k">return</span> <span class="p">[</span><span class="n">topi_compute</span><span class="p">(</span><span class="n">inputs</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">attrs</span><span class="o">.</span><span class="n">axis</span><span class="p">,</span> <span class="n">attrs</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">attrs</span><span class="o">.</span><span class="n">exclusive</span><span class="p">)]</span>

    <span class="k">return</span> <span class="n">_compute_scanop</span>


<span class="nd">@override_native_generic_func</span><span class="p">(</span><span class="s2">&quot;cumsum_strategy&quot;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">cumsum_strategy</span><span class="p">(</span><span class="n">attrs</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">out_type</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;cumsum generic strategy&quot;&quot;&quot;</span>
    <span class="n">strategy</span> <span class="o">=</span> <span class="n">_op</span><span class="o">.</span><span class="n">OpStrategy</span><span class="p">()</span>
    <span class="n">strategy</span><span class="o">.</span><span class="n">add_implementation</span><span class="p">(</span>
        <span class="n">wrap_compute_scanop</span><span class="p">(</span><span class="n">topi</span><span class="o">.</span><span class="n">cumsum</span><span class="p">),</span>
        <span class="n">wrap_topi_schedule</span><span class="p">(</span><span class="n">topi</span><span class="o">.</span><span class="n">generic</span><span class="o">.</span><span class="n">schedule_extern</span><span class="p">),</span>
        <span class="n">name</span><span class="o">=</span><span class="s2">&quot;cumsum.generic&quot;</span><span class="p">,</span>
    <span class="p">)</span>
    <span class="k">return</span> <span class="n">strategy</span>


<span class="nd">@override_native_generic_func</span><span class="p">(</span><span class="s2">&quot;cumprod_strategy&quot;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">cumprod_strategy</span><span class="p">(</span><span class="n">attrs</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">out_type</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;cumprod generic strategy&quot;&quot;&quot;</span>
    <span class="n">strategy</span> <span class="o">=</span> <span class="n">_op</span><span class="o">.</span><span class="n">OpStrategy</span><span class="p">()</span>
    <span class="n">strategy</span><span class="o">.</span><span class="n">add_implementation</span><span class="p">(</span>
        <span class="n">wrap_compute_scanop</span><span class="p">(</span><span class="n">topi</span><span class="o">.</span><span class="n">cumprod</span><span class="p">),</span>
        <span class="n">wrap_topi_schedule</span><span class="p">(</span><span class="n">topi</span><span class="o">.</span><span class="n">generic</span><span class="o">.</span><span class="n">schedule_extern</span><span class="p">),</span>
        <span class="n">name</span><span class="o">=</span><span class="s2">&quot;cumprod.generic&quot;</span><span class="p">,</span>
    <span class="p">)</span>
    <span class="k">return</span> <span class="n">strategy</span>

<span class="nd">@cumsum_strategy.register</span><span class="p">([</span><span class="s2">&quot;cuda&quot;</span><span class="p">,</span> <span class="s2">&quot;gpu&quot;</span><span class="p">])</span>
<span class="k">def</span> <span class="nf">cumsum_strategy_cuda</span><span class="p">(</span><span class="n">attrs</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">out_type</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;cumsum cuda strategy&quot;&quot;&quot;</span>
    <span class="n">strategy</span> <span class="o">=</span> <span class="n">_op</span><span class="o">.</span><span class="n">OpStrategy</span><span class="p">()</span>
    <span class="n">strategy</span><span class="o">.</span><span class="n">add_implementation</span><span class="p">(</span>
        <span class="n">wrap_compute_scanop</span><span class="p">(</span><span class="n">topi</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">cumsum</span><span class="p">),</span>
        <span class="n">wrap_topi_schedule</span><span class="p">(</span><span class="n">topi</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">schedule_scan</span><span class="p">),</span>
        <span class="n">name</span><span class="o">=</span><span class="s2">&quot;cumsum.cuda&quot;</span><span class="p">,</span>
    <span class="p">)</span>
    <span class="k">return</span> <span class="n">strategy</span>


<span class="nd">@cumprod_strategy.register</span><span class="p">([</span><span class="s2">&quot;cuda&quot;</span><span class="p">,</span> <span class="s2">&quot;gpu&quot;</span><span class="p">])</span>
<span class="k">def</span> <span class="nf">cumprod_strategy_cuda</span><span class="p">(</span><span class="n">attrs</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">out_type</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;cumprod cuda strategy&quot;&quot;&quot;</span>
    <span class="n">strategy</span> <span class="o">=</span> <span class="n">_op</span><span class="o">.</span><span class="n">OpStrategy</span><span class="p">()</span>
    <span class="n">strategy</span><span class="o">.</span><span class="n">add_implementation</span><span class="p">(</span>
        <span class="n">wrap_compute_scanop</span><span class="p">(</span><span class="n">topi</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">cumprod</span><span class="p">),</span>
        <span class="n">wrap_topi_schedule</span><span class="p">(</span><span class="n">topi</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">schedule_scan</span><span class="p">),</span>
        <span class="n">name</span><span class="o">=</span><span class="s2">&quot;cumprod.cuda&quot;</span><span class="p">,</span>
    <span class="p">)</span>
    <span class="k">return</span> <span class="n">strategy</span>
</pre></div>
</div>
<p>Where in each strategy we define the compute we wrote and the schedule to use within <code class="docutils literal notranslate"><span class="pre">add_implementation()</span></code>.
We finally link the strategy and compute with the defined relay operator in <code class="docutils literal notranslate"><span class="pre">python/tvm/relay/op/_transform.py</span></code>:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="c1"># cumsum</span>
<span class="nd">@_reg.register_compute</span><span class="p">(</span><span class="s2">&quot;cumsum&quot;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">compute_cumsum</span><span class="p">(</span><span class="n">attrs</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">output_type</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Compute definition of cumsum&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="p">[</span><span class="n">topi</span><span class="o">.</span><span class="n">cumsum</span><span class="p">(</span><span class="n">inputs</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">attrs</span><span class="o">.</span><span class="n">axis</span><span class="p">,</span> <span class="n">attrs</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">attrs</span><span class="o">.</span><span class="n">exclusive</span><span class="p">)]</span>


<span class="n">_reg</span><span class="o">.</span><span class="n">register_strategy</span><span class="p">(</span><span class="s2">&quot;cumsum&quot;</span><span class="p">,</span> <span class="n">strategy</span><span class="o">.</span><span class="n">cumsum_strategy</span><span class="p">)</span>
<span class="n">_reg</span><span class="o">.</span><span class="n">register_shape_func</span><span class="p">(</span><span class="s2">&quot;cumsum&quot;</span><span class="p">,</span> <span class="bp">False</span><span class="p">,</span> <span class="n">elemwise_shape_func</span><span class="p">)</span>

<span class="c1"># cumprod</span>
<span class="nd">@_reg.register_compute</span><span class="p">(</span><span class="s2">&quot;cumprod&quot;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">compute_cumprod</span><span class="p">(</span><span class="n">attrs</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">output_type</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Compute definition of cumprod&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="p">[</span><span class="n">topi</span><span class="o">.</span><span class="n">cumprod</span><span class="p">(</span><span class="n">inputs</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">attrs</span><span class="o">.</span><span class="n">axis</span><span class="p">,</span> <span class="n">attrs</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">attrs</span><span class="o">.</span><span class="n">exclusive</span><span class="p">)]</span>


<span class="n">_reg</span><span class="o">.</span><span class="n">register_strategy</span><span class="p">(</span><span class="s2">&quot;cumprod&quot;</span><span class="p">,</span> <span class="n">strategy</span><span class="o">.</span><span class="n">cumprod_strategy</span><span class="p">)</span>
<span class="n">_reg</span><span class="o">.</span><span class="n">register_shape_func</span><span class="p">(</span><span class="s2">&quot;cumprod&quot;</span><span class="p">,</span> <span class="bp">False</span><span class="p">,</span> <span class="n">elemwise_shape_func</span><span class="p">)</span>
</pre></div>
</div>
<p>The shape functions are used for determining output shape given a dynamically shaped tensor. In this
case we tell TVM the output shape will be the same as the input shape.</p>
</div>
<div class="section" id="creating-a-relay-call-node-and-exposing-a-python-hook">
<h2>6. Creating a Relay Call Node and Exposing a Python Hook<a class="headerlink" href="#creating-a-relay-call-node-and-exposing-a-python-hook" title="永久链接至标题">¶</a></h2>
<p>We now have a working operation and now just need to properly call it
via a Relay Call Node. This step requires simply writing a function that takes
the arguments to the operator (as Relay expressions) and
returning a call node to the operator (i.e., the node that
should be placed into the Relay AST where the call to the
operator is intended).</p>
<p>At present call attributes and type arguments (the last two fields)
are not supported, so it suffices to use <code class="docutils literal notranslate"><span class="pre">Op::Get</span></code> to fetch
the operator’s information from the operator registry and pass in
the arguments to the call node, as below. In <code class="docutils literal notranslate"><span class="pre">src/relay/op/tensor/transform.cc</span></code>:</p>
<div class="highlight-c++ notranslate"><div class="highlight"><pre><span></span><span class="n">Expr</span> <span class="nf">MakeCumsum</span><span class="p">(</span><span class="n">Expr</span> <span class="n">data</span><span class="p">,</span> <span class="n">Integer</span> <span class="n">axis</span><span class="p">,</span> <span class="n">DataType</span> <span class="n">dtype</span><span class="p">,</span> <span class="n">Bool</span> <span class="n">exclusive</span><span class="p">)</span> <span class="p">{</span>
    <span class="k">auto</span> <span class="n">attrs</span> <span class="o">=</span> <span class="n">make_object</span><span class="o">&lt;</span><span class="n">ScanopAttrs</span><span class="o">&gt;</span><span class="p">();</span>
    <span class="n">attrs</span><span class="o">-&gt;</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span><span class="p">;</span>
    <span class="n">attrs</span><span class="o">-&gt;</span><span class="n">axis</span> <span class="o">=</span> <span class="n">axis</span><span class="p">;</span>
    <span class="n">attrs</span><span class="o">-&gt;</span><span class="n">exclusive</span> <span class="o">=</span> <span class="n">exclusive</span><span class="p">;</span>
    <span class="k">static</span> <span class="k">const</span> <span class="n">Op</span><span class="o">&amp;</span> <span class="n">op</span> <span class="o">=</span> <span class="n">Op</span><span class="o">::</span><span class="n">Get</span><span class="p">(</span><span class="s">&quot;cumsum&quot;</span><span class="p">);</span>
    <span class="k">return</span> <span class="n">Call</span><span class="p">(</span><span class="n">op</span><span class="p">,</span> <span class="p">{</span><span class="n">data</span><span class="p">},</span> <span class="n">Attrs</span><span class="p">(</span><span class="n">attrs</span><span class="p">),</span> <span class="p">{});</span>
<span class="p">}</span>

<span class="n">TVM_REGISTER_GLOBAL</span><span class="p">(</span><span class="s">&quot;relay.op._make.cumsum&quot;</span><span class="p">).</span><span class="n">set_body_typed</span><span class="p">(</span><span class="n">MakeCumsum</span><span class="p">);</span>

<span class="n">Expr</span> <span class="nf">MakeCumprod</span><span class="p">(</span><span class="n">Expr</span> <span class="n">data</span><span class="p">,</span> <span class="n">Integer</span> <span class="n">axis</span><span class="p">,</span> <span class="n">DataType</span> <span class="n">dtype</span><span class="p">,</span> <span class="n">Bool</span> <span class="n">exclusive</span><span class="p">)</span> <span class="p">{</span>
    <span class="k">auto</span> <span class="n">attrs</span> <span class="o">=</span> <span class="n">make_object</span><span class="o">&lt;</span><span class="n">ScanopAttrs</span><span class="o">&gt;</span><span class="p">();</span>
    <span class="n">attrs</span><span class="o">-&gt;</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span><span class="p">;</span>
    <span class="n">attrs</span><span class="o">-&gt;</span><span class="n">axis</span> <span class="o">=</span> <span class="n">axis</span><span class="p">;</span>
    <span class="n">attrs</span><span class="o">-&gt;</span><span class="n">exclusive</span> <span class="o">=</span> <span class="n">exclusive</span><span class="p">;</span>
    <span class="k">static</span> <span class="k">const</span> <span class="n">Op</span><span class="o">&amp;</span> <span class="n">op</span> <span class="o">=</span> <span class="n">Op</span><span class="o">::</span><span class="n">Get</span><span class="p">(</span><span class="s">&quot;cumprod&quot;</span><span class="p">);</span>
    <span class="k">return</span> <span class="n">Call</span><span class="p">(</span><span class="n">op</span><span class="p">,</span> <span class="p">{</span><span class="n">data</span><span class="p">},</span> <span class="n">Attrs</span><span class="p">(</span><span class="n">attrs</span><span class="p">),</span> <span class="p">{});</span>
<span class="p">}</span>

<span class="n">TVM_REGISTER_GLOBAL</span><span class="p">(</span><span class="s">&quot;relay.op._make.cumsum&quot;</span><span class="p">).</span><span class="n">set_body_typed</span><span class="p">(</span><span class="n">MakeCumprod</span><span class="p">);</span>
</pre></div>
</div>
<p>Where <code class="docutils literal notranslate"><span class="pre">TVM_REGISTER_GLOBAL</span></code> exposes the <code class="docutils literal notranslate"><span class="pre">MakeCumsum</span></code> and <code class="docutils literal notranslate"><span class="pre">MakeCumprod</span></code> functions
in Python via <code class="docutils literal notranslate"><span class="pre">relay.op._make.cumsum(...)</span></code> and <code class="docutils literal notranslate"><span class="pre">relay.op._make.cumsum(...)</span></code>.</p>
</div>
<div class="section" id="including-a-cleaner-python-api-hook">
<h2>7. Including a Cleaner Python API Hook<a class="headerlink" href="#including-a-cleaner-python-api-hook" title="永久链接至标题">¶</a></h2>
<p>It is generally the convention in Relay, that functions exported
through <code class="docutils literal notranslate"><span class="pre">TVM_REGISTER_GLOBAL</span></code> should be wrapped in a separate
Python function rather than called directly in Python. For our
operators we expose this cleaner interface in <code class="docutils literal notranslate"><span class="pre">python/tvm/relay/op/transform.py</span></code></p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">cumsum</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="bp">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">None</span><span class="p">,</span> <span class="n">exclusive</span><span class="o">=</span><span class="bp">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">_make</span><span class="o">.</span><span class="n">cumsum</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">axis</span><span class="p">,</span> <span class="n">dtype</span><span class="p">,</span> <span class="n">exclusive</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">cumprod</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="bp">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">None</span><span class="p">,</span> <span class="n">exclusive</span><span class="o">=</span><span class="bp">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">_make</span><span class="o">.</span><span class="n">cumprod</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">axis</span><span class="p">,</span> <span class="n">dtype</span><span class="p">,</span> <span class="n">exclusive</span><span class="p">)</span>
</pre></div>
</div>
<p>Note that these Python wrappers might also be good opportunities to
provide an easier interface to the operator. For example, the
<code class="docutils literal notranslate"><span class="pre">concat</span></code> operator is registered as taking only one operator,
namely a tuple with the tensors to be concatenated, but the Python
wrapper takes the tensors as arguments and combines them into a tuple
before producing the call node:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">concat</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Concatenate the input tensors along the zero axis.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    args: list of Tensor</span>

<span class="sd">    Returns</span>
<span class="sd">    -------</span>
<span class="sd">    tensor: The concatenated tensor.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">tup</span> <span class="o">=</span> <span class="n">Tuple</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">args</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">_make</span><span class="o">.</span><span class="n">concat</span><span class="p">(</span><span class="n">tup</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="section" id="writing-unit-tests">
<h2>8. Writing Unit Tests!<a class="headerlink" href="#writing-unit-tests" title="永久链接至标题">¶</a></h2>
<p>This is self explanatory! Some example unit tests can be found in
<a class="reference external" href="https://github.com/apache/tvm/blob/main/tests/python/relay/test_op_level3.py">tests/python/relay/test_op_level3.py</a> for our cumulative sum
and product operators.</p>
</div>
<div class="section" id="other-topics">
<h2>Other Topics<a class="headerlink" href="#other-topics" title="永久链接至标题">¶</a></h2>
<div class="section" id="gradient-operators">
<h3>Gradient Operators<a class="headerlink" href="#gradient-operators" title="永久链接至标题">¶</a></h3>
<p>Gradient operators are important for writing differentiable programs in
Relay. While it is the case that Relay’s autodiff algorithm can differentiate
first-class language constructs, operators are opaque. Because Relay can’t
look into the implementation, an explicit differentiation rule must be
provided.</p>
<p>Both Python and C++ can be used to write gradient operators, but we focus our
examples on Python, as it is more commonly used.</p>
</div>
<div class="section" id="adding-a-gradient-in-python">
<h3>Adding a Gradient in Python<a class="headerlink" href="#adding-a-gradient-in-python" title="永久链接至标题">¶</a></h3>
<p>A collection of Python gradient operators can be found in
<code class="docutils literal notranslate"><span class="pre">python/tvm/relay/op/_tensor_grad.py</span></code>. We will walk through two
representative examples: <code class="docutils literal notranslate"><span class="pre">sigmoid</span></code> and <code class="docutils literal notranslate"><span class="pre">multiply</span></code>.</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="nd">@register_gradient</span><span class="p">(</span><span class="s2">&quot;sigmoid&quot;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">sigmoid_grad</span><span class="p">(</span><span class="n">orig</span><span class="p">,</span> <span class="n">grad</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Returns [grad * sigmoid(x) * (1 - sigmoid(x))].&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="p">[</span><span class="n">grad</span> <span class="o">*</span> <span class="n">orig</span> <span class="o">*</span> <span class="p">(</span><span class="n">ones_like</span><span class="p">(</span><span class="n">orig</span><span class="p">)</span> <span class="o">-</span> <span class="n">orig</span><span class="p">)]</span>
</pre></div>
</div>
<p>The inputs here are the original operator <code class="docutils literal notranslate"><span class="pre">orig</span></code> and a gradient <code class="docutils literal notranslate"><span class="pre">grad</span></code> to
accumulate into. What we return is a list, where the element at the i’th
index is the derivative of the operator with respect to the operator’s i’th
input. In general, the gradient will return a list with as many elements as
there are inputs to the base operator.</p>
<p>Before we further analyze this definition, first we should recall the
derivative of the sigmoid function: <span class="math notranslate nohighlight">\(\frac{\partial \sigma}{\partial x}
= \sigma(x)(1 - \sigma(x))\)</span>. The definition above looks similar to the
mathematical definition, but there is one important addition, which we
describe below.</p>
<p>The term <code class="docutils literal notranslate"><span class="pre">orig</span> <span class="pre">*</span> <span class="pre">(ones_like(orig)</span> <span class="pre">-</span> <span class="pre">orig)</span></code> directly matches the derivative,
because <code class="docutils literal notranslate"><span class="pre">orig</span></code> here is the sigmoid function, but we’re not just interested
in how to compute the gradient of this function. We’re interested in
composing this gradient with other gradients, so we can accumulate the
gradient across an entire program. This is where the <code class="docutils literal notranslate"><span class="pre">grad</span></code> term comes in.
In the expression <code class="docutils literal notranslate"><span class="pre">grad</span> <span class="pre">*</span> <span class="pre">orig</span> <span class="pre">*</span> <span class="pre">(ones_like(orig)</span> <span class="pre">-</span> <span class="pre">orig)</span></code>, multiplying by
<code class="docutils literal notranslate"><span class="pre">grad</span></code> specifies how to compose the derivative with the gradient thus far.</p>
<p>Now, we consider <code class="docutils literal notranslate"><span class="pre">multiply</span></code>, a slightly more interesting example:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="nd">@register_gradient</span><span class="p">(</span><span class="s2">&quot;multiply&quot;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">multiply_grad</span><span class="p">(</span><span class="n">orig</span><span class="p">,</span> <span class="n">grad</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Returns [grad * y, grad * x]&quot;&quot;&quot;</span>
    <span class="n">x</span><span class="p">,</span> <span class="n">y</span> <span class="o">=</span> <span class="n">orig</span><span class="o">.</span><span class="n">args</span>
    <span class="k">return</span> <span class="p">[</span><span class="n">collapse_sum_like</span><span class="p">(</span><span class="n">grad</span> <span class="o">*</span> <span class="n">y</span><span class="p">,</span> <span class="n">x</span><span class="p">),</span>
            <span class="n">collapse_sum_like</span><span class="p">(</span><span class="n">grad</span> <span class="o">*</span> <span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">)]</span>
</pre></div>
</div>
<p>In this example, there are two elements in the returned list, because
<code class="docutils literal notranslate"><span class="pre">multiply</span></code> is a binary operator. And to recall, if <span class="math notranslate nohighlight">\(f(x, y) = xy\)</span>, the
partial derivatives are <span class="math notranslate nohighlight">\(\frac{\partial f}{\partial x} = y\)</span> and
<span class="math notranslate nohighlight">\(\frac{\partial f}{\partial y} = x\)</span>.</p>
<p>There is one required step for <code class="docutils literal notranslate"><span class="pre">multiply</span></code> that is not required for
<code class="docutils literal notranslate"><span class="pre">sigmoid</span></code>, because <code class="docutils literal notranslate"><span class="pre">multiply</span></code> has broadcasting semantics. Since the shape
of <code class="docutils literal notranslate"><span class="pre">grad</span></code> might not match the shape of the inputs, we use
<code class="docutils literal notranslate"><span class="pre">collapse_sum_like</span></code> to take the contents of the <code class="docutils literal notranslate"><span class="pre">grad</span> <span class="pre">*</span> <span class="pre">&lt;var&gt;</span></code> terms and
make the shape match the shape of the input we’re differentiating with
respect to.</p>
</div>
<div class="section" id="adding-a-gradient-in-c">
<h3>Adding a Gradient in C++<a class="headerlink" href="#adding-a-gradient-in-c" title="永久链接至标题">¶</a></h3>
<p>Adding a gradient in C++ is similar to adding one in Python, but the
interface for registering is slightly different.</p>
<p>First, make sure <code class="docutils literal notranslate"><span class="pre">src/relay/transforms/pattern_utils.h</span></code> is included. It provides
helper functions for creating nodes in the Relay AST. Then, define the
gradient in a similar fashion as in the Python example:</p>
<div class="highlight-c notranslate"><div class="highlight"><pre><span></span><span class="n">tvm</span><span class="o">::</span><span class="n">Array</span><span class="o">&lt;</span><span class="n">Expr</span><span class="o">&gt;</span> <span class="n">MultiplyGrad</span><span class="p">(</span><span class="k">const</span> <span class="n">Expr</span><span class="o">&amp;</span> <span class="n">orig_call</span><span class="p">,</span> <span class="k">const</span> <span class="n">Expr</span><span class="o">&amp;</span> <span class="n">output_grad</span><span class="p">)</span> <span class="p">{</span>
    <span class="k">const</span> <span class="n">Call</span><span class="o">&amp;</span> <span class="n">call</span> <span class="o">=</span> <span class="n">orig_call</span><span class="p">.</span><span class="n">Downcast</span><span class="o">&lt;</span><span class="n">Call</span><span class="o">&gt;</span><span class="p">();</span>
    <span class="k">return</span> <span class="p">{</span> <span class="n">CollapseSumLike</span><span class="p">(</span><span class="n">Multiply</span><span class="p">(</span><span class="n">output_grad</span><span class="p">,</span> <span class="n">call</span><span class="p">.</span><span class="n">args</span><span class="p">[</span><span class="mi">1</span><span class="p">]),</span> <span class="n">call</span><span class="p">.</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]),</span>
             <span class="n">CollapseSumLike</span><span class="p">(</span><span class="n">Multiply</span><span class="p">(</span><span class="n">output_grad</span><span class="p">,</span> <span class="n">call</span><span class="p">.</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]),</span> <span class="n">call</span><span class="p">.</span><span class="n">args</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="p">};</span>
<span class="p">}</span>
</pre></div>
</div>
<p>Notice that in C++ we can’t use the same operator overloading that we have in
Python, and we need to downcast, so the implementation is more verbose. Even
so, we can easily verify that this definition mirrors the earlier example in
Python.</p>
<p>Now, instead of using a Python decorator, we need to tack a <code class="docutils literal notranslate"><span class="pre">set_attr</span></code> call
for “FPrimalGradient” onto the end of the base operator’s registration, in
order to register the gradient.</p>
<div class="highlight-c notranslate"><div class="highlight"><pre><span></span><span class="n">RELAY_REGISTER_OP</span><span class="p">(</span><span class="s">&quot;multiply&quot;</span><span class="p">)</span>
    <span class="c1">// ...</span>
    <span class="c1">// Set other attributes</span>
    <span class="c1">// ...</span>
    <span class="p">.</span><span class="n">set_attr</span><span class="o">&lt;</span><span class="n">FPrimalGradient</span><span class="o">&gt;</span><span class="p">(</span><span class="s">&quot;FPrimalGradient&quot;</span><span class="p">,</span> <span class="n">MultiplyGrad</span><span class="p">);</span>
</pre></div>
</div>
</div>
</div>
</div>


           </div>
           
          </div>
          

<footer>

    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="relay_add_pass.html" class="btn btn-neutral float-right" title="Adding a Compiler Pass to Relay" accesskey="n" rel="next">下一个 <span class="fa fa-arrow-circle-right"></span></a>
      
      
        <a href="how_to.html" class="btn btn-neutral float-left" title="开发者指南" accesskey="p" rel="prev"><span class="fa fa-arrow-circle-left"></span> 上一个</a>
      
    </div>

<div id="button" class="backtop"><img src="../../_static//img/right.svg" alt="backtop"/> </div>
<section class="footerSec">
    <div class="footerHeader">
      <ul class="d-flex align-md-items-center justify-content-between flex-column flex-md-row">
        <li class="copywrite d-flex align-items-center">
          <h5 id="copy-right-info">© 2020 Apache Software Foundation | All right reserved</h5>
        </li>
      </ul>

    </div>

    <ul>
      <li class="footernote">Copyright © 2020 The Apache Software Foundation. Apache TVM, Apache, the Apache feather, and the Apache TVM project logo are either trademarks or registered trademarks of the Apache Software Foundation.</li>
    </ul>

</section>
</footer>
        </div>
      </div>

    </section>

  </div>
  

    <script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.12.9/umd/popper.min.js" integrity="sha384-ApNbgh9B+Y1QKtv3Rn7W3mgPxhU9K/ScQsAP7hUibX39j7fakFPskvXusvfa0b4Q" crossorigin="anonymous"></script>
    <script src="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/js/bootstrap.min.js" integrity="sha384-JZR6Spejh4U02d8jOt6vLEHfe/JQGiRRSQQxSfFWpi1MquVdAyjUar5+76PVCmYl" crossorigin="anonymous"></script>

  </body>
  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
    <!-- Theme Analytics -->
    <script>
    (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
      (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o),
      m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
    })(window,document,'script','https://www.google-analytics.com/analytics.js','ga');

    ga('create', 'UA-75982049-2', 'auto');
    ga('send', 'pageview');
    </script>

    
   

</body>
</html>